import cmd
import re

import paramiko

# 字符串末尾换行符正则
PATTERN_LAST_LINEBREAK = re.compile('[\n\r]$')
# 路径末尾正斜杠正则
PATTERN_LAST_PATH_SEP = re.compile('[/]$')
# 编码列表
SUPPORT_ENCODINGS = ['UTF-8', 'GBK', 'GB2312', 'ISO-8859-1']


class SSHTerminal(cmd.Cmd):
    default_wd = ''
    # 当前会话所在目录
    current_wd = ''
    encoding = 'utf-8'
    # 历史命令，只记录执行成功的命令
    history = []

    def __init__(self, host, port, user, password):
        super(SSHTerminal, self).__init__()
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.ssh_client = paramiko.SSHClient()

        self.default_wd = '/root' if user == 'root' else ('/home/' + user)
        self.set_curren_wd(self.default_wd)

    def preloop(self):
        try:
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            self.ssh_client.connect(self.host, self.port, self.user, self.password, timeout=7)
        except Exception as e:
            print(str(e))
            exit(0)

    def exec_command(self, command):
        # print('command ==> ', command)
        stdin, stdout, stderr = self.ssh_client.exec_command(command=command)
        out_res = str(stdout.read(), encoding=self.encoding)
        err_res = str(stderr.read(), encoding=self.encoding)
        return re.sub(PATTERN_LAST_LINEBREAK, '', out_res), re.sub(PATTERN_LAST_LINEBREAK, '', err_res)

    def set_curren_wd(self, path):
        """ 设置当前shell所在目录 """
        if path:
            if path.startswith('~'):
                path = re.sub('^~', self.default_wd, path)
            if path.startswith('/') or path.startswith(self.default_wd):
                # 目录从home目录或者根目录开始，直接赋值
                self.current_wd = path
            else:
                paths = re.sub(PATTERN_LAST_PATH_SEP, '', path).split('/')
                # 路径中有返回上层..符号
                if paths.count('..') > 0 or paths.count('.') > 0:
                    # 当前目录在根目录，然后执行了cd ../xxx命令，直接返回
                    if path.startswith('..') and self.current_wd == '/':
                        return
                    cur_paths = re.sub(PATTERN_LAST_PATH_SEP, '', self.current_wd).split('/')
                    for p in paths:
                        if p == '..':
                            cur_paths.pop()
                        elif p == '.':
                            continue
                        else:
                            cur_paths.append(p)
                    self.current_wd = '/'.join(cur_paths) or '/'
                else:
                    self.current_wd += path if self.current_wd and self.current_wd.endswith('/') else ('/' + path)

            # 设置提示字符中显示的路径
            if self.current_wd == '/' or self.current_wd == self.default_wd:
                # 在用户目录下,显示波浪线
                cd = '~' if self.current_wd == self.default_wd else self.current_wd
            else:
                p = re.sub(PATTERN_LAST_PATH_SEP, '', self.current_wd)
                cd = p[p.rindex('/') + 1:]
            self.prompt = '[%s@%s %s]# ' % (self.user, self.host, cd)

    def set_encoding(self, encoding):
        """ 设置shell编码，默认为UTF-8 """
        self.encoding = encoding or 'utf-8'

    def default(self, line):
        try:
            command, args, line = self.parseline(line)
            if not command:
                return
            if command == 'cd' and not args:
                # 只输入cd命令默认切换到用户目录
                line += ' ' + self.default_wd
            result, error = self.exec_command('cd %s && %s' % (self.current_wd, line))
            if error:
                print(error)
            else:
                if command == 'cd':
                    self.set_curren_wd((args.strip() or self.default_wd) if args else self.default_wd)
                else:
                    result and print(result)
                self.history.append(line)
        except Exception as e:
            print(e)

    def do_curwd(self, args):
        print(self.current_wd)

    def do_history(self, args):
        """ 显示历史命令 """
        for his in self.history:
            print(his)

    def do_se(self, args):
        """ 设置编码 """
        encoding = args.upper()
        if encoding in SUPPORT_ENCODINGS:
            self.set_encoding(encoding)
        else:
            print('不支持的编码类型[%s]' % args)

    def do_exit(self, args):
        print('Bye')
        self.ssh_client.close()
        return True

    def do_quit(self, args):
        return self.do_exit(args)

    def do_bye(self, args):
        return self.do_exit(args)

    def do_help(self, arg):
        print('终端说明文档：')
        print('命令：')
        print('  se              设置Shell编码类型，例：se utf-8')
        print('  history         查看历史输入命令，只会显示远程成功执行的命令')
        print('  exit/quit/bye   退出终端')
        print('  help/?          查看帮助文档')
        print('  其他Linux命令，暂不支持有交互的命令')
